// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.

using System;
using System.Threading;
using System.Threading.Tasks;

#nullable enable
namespace Ignitor;

internal class CancellableOperation<TResult>
{
    public CancellableOperation(TimeSpan? timeout, CancellationToken cancellationToken)
    {
        Timeout = timeout;

        Completion = new TaskCompletionSource<TResult>(TaskCreationOptions.RunContinuationsAsynchronously);
        Completion.Task.ContinueWith(
            (task, state) =>
            {
                var operation = (CancellableOperation<TResult>)state!;
                operation.Dispose();
            },
            this,
            cancellationToken,
            TaskContinuationOptions.ExecuteSynchronously, // We need to execute synchronously to clean-up before anything else continues
            TaskScheduler.Default);

        Cancellation = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);

        if (Timeout != null && Timeout != System.Threading.Timeout.InfiniteTimeSpan && Timeout != TimeSpan.MaxValue)
        {
            Cancellation.CancelAfter(Timeout.Value);
        }

        CancellationRegistration = Cancellation.Token.Register(
            (self) =>
            {
                var operation = (CancellableOperation<TResult>)self!;

                if (cancellationToken.IsCancellationRequested)
                {
                        // The operation was externally canceled before it timed out.
                        Dispose();
                    return;
                }

                operation.Completion.TrySetException(new TimeoutException($"The operation timed out after {Timeout}."));
                operation.Cancellation?.Dispose();
                operation.CancellationRegistration.Dispose();
            },
            this);
    }

    public TimeSpan? Timeout { get; }

    public TaskCompletionSource<TResult> Completion { get; }

    public CancellationTokenSource Cancellation { get; }

    public CancellationTokenRegistration CancellationRegistration { get; }

    public bool Disposed { get; private set; }

    private void Dispose()
    {
        if (Disposed)
        {
            return;
        }

        Disposed = true;
        Completion.TrySetCanceled(Cancellation.Token);
        Cancellation.Dispose();
        CancellationRegistration.Dispose();
    }
}
#nullable restore
